import numpy as np
import faiss
import dashscope
import openai

# ------ 配置 API Key ------ #
dashscope.api_key = "你的DashScope Key"
openai.api_key = "你的OpenAI Key"  # 用于查询改写

# ------ 1. 查询改写（使用 GPT） ------ #
def rewrite_query(query: str, n=3):
    prompt = f"请将用户问题“{query}”改写为{n}种意思相近但表达不同的语句，用于语义搜索。"
    completion = openai.ChatCompletion.create(
        model="gpt-3.5-turbo",
        messages=[{"role": "user", "content": prompt}],
        temperature=0.3,
    )
    lines = completion['choices'][0]['message']['content'].split('\n')
    return [query] + [l.strip("0123456789.-、 ") for l in lines if l.strip()]

# ------ 2. 获取 DashScope embedding 向量 ------ #
def get_dashscope_embedding(text: str) -> np.ndarray:
    response = dashscope.TextEmbedding.call(
        model="text-embedding-v2",
        input=text,
    )
    return np.array(response.output["embeddings"][0]["embedding"], dtype=np.float32)

# ------ 3. 加载 FAISS 向量库 & 文本映射 ------ #
index = faiss.read_index("my_index.index")  # 你的向量库路径
with open("doc_texts.txt", "r", encoding="utf-8") as f:
    doc_texts = [line.strip() for line in f.readlines()]

# ------ 4. 搜索流程 ------ #
query = "法务合规管理"
expanded_queries = rewrite_query(query)

results = []
for q in expanded_queries:
    emb = get_dashscope_embedding(q)
    D, I = index.search(emb.reshape(1, -1), k=5)
    for dist, idx in zip(D[0], I[0]):
        results.append((dist, doc_texts[idx]))

# ------ 5. 去重 + 排序 ------ #
unique = {}
for dist, text in results:
    if text not in unique or dist < unique[text]:
        unique[text] = dist

sorted_results = sorted(unique.items(), key=lambda x: x[1])

# ------ 6. 输出结果 ------ #
for i, (text, dist) in enumerate(sorted_results[:10]):
    print(f"\n【Top {i+1} | 距离: {dist:.4f}】\n{text}")
